# -- coding: utf-8 --
"""
Tensorflow框架内网络调用 或 自定义网络结构
"""   
from __future__ import absolute_import #在 3.0 以前的旧版本中启用相对导入等特性所必须的 future 语句
from __future__ import division
import os.path
import re
import tensorflow.python.platform
import tensorflow as tf
from prostate_input import inputPipeLine
import pdb
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim import nets
vgg=nets.vgg
vgg16 = nets.vgg.vgg_16

inception_v3 = nets.inception.inception_v3
inception_v3_arg_scope = nets.inception.inception_v3_arg_scope
inception_resnet_v2 = nets.inception.inception_resnet_v2
inception_resnet_v2_arg_scope = nets.inception.inception_resnet_v2_arg_scope

resnet_v2 = nets.resnet_v2
resnet50 = nets.resnet_v2.resnet_v2_50
resnet101 = nets.resnet_v2.resnet_v2_101
resnet152 = nets.resnet_v2.resnet_v2_152

FLAGS = tf.app.flags.FLAGS

def _activation_summary(x):	#summary记录每层计算的tensor（特征图）和稀疏度
	tensor_name = re.sub('%s_[0-9]*/' % TOWER_NAME, '', x.op.name)
	tf.summary.histogram(tensor_name + '/activations', x)
	tf.summary.scalar(tensor_name + '/sparsity', tf.nn.zero_fraction(x))

def _activation_summaries(endpoints):
	with tf.name_scope('summaries'):
		for act in endpoints.values():
			_activation_summary(act)

def _variable_on_cpu(name, shape, initializer):
	with tf.device('/cpu:0'):
		var = tf.get_variable(name, shape, initializer=initializer)
	return var
  
def _variable_with_weight_decay(name, shape, stddev, wd):#增加权重衰减项
	var = _variable_on_cpu(name, shape,
						 tf.truncated_normal_initializer(stddev=stddev))
	if wd:
		weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') #multiply各元素对应相乘
		tf.add_to_collection('losses', weight_decay)
	return var

def vgg16_inference(image_batch, weight_decay=0.0004,scope=None, is_training=True):
	with slim.arg_scope(vgg.vgg_arg_scope(weight_decay=weight_decay)):
		logits, end_points = vgg16(
			image_batch,
			num_classes=2,
			is_training=is_training
			)
	_activation_summaries(end_points)
	return logits, end_points 

def inception_v3_inference(image_batch, weight_decay=0.0004,scope=None, is_training=None, reuse = None):
	with slim.arg_scope(inception_v3_arg_scope(weight_decay=weight_decay)):
		logits, end_points = inception_v3(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	_activation_summaries(end_points)
	auxiliary_logits = end_points['AuxLogits']
	return logits, end_points #logits, auxiliary_logits, end_points['Predictions']


def inception_v4_inference(image_batch, weight_decay=0.0004,scope=None, is_training=None, reuse = None):
	with slim.arg_scope(inception_v4_arg_scope(weight_decay=weight_decay)):
		logits, end_points = inception_v4(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	_activation_summaries(end_points)
	auxiliary_logits = end_points['AuxLogits']
	return logits, end_points #logits, auxiliary_logits, end_points['Predictions']

def inception_resnet_v2_inference(image_batch, weight_decay=0.0004,scope=None, is_training=None, reuse = None):
	with slim.arg_scope(inception_resnet_v2_arg_scope(weight_decay=weight_decay)):
		logits, end_points = inception_resnet_v2(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	_activation_summaries(end_points)
	auxiliary_logits = end_points['AuxLogits']
	return logits, end_points 

def resnet152_inference(image_batch, weight_decay=0.0004,scope=None, is_training=True, reuse = None):
	with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
		net, end_points = resnet152(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	logits = tf.squeeze(net)
	_activation_summaries(end_points)
	return logits, end_points 
	#return net, end_points['predictions']

def resnet101_inference(image_batch, weight_decay=0.0004,scope=None, is_training=True, reuse = None):
	with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
		net, end_points = resnet101(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	logits = tf.squeeze(net)
	# _activation_summaries(end_points)
	return logits, end_points

def resnet50_inference(image_batch, weight_decay=0.0004,scope=None, is_training=True, reuse = None):
	with slim.arg_scope(resnet_v2.resnet_arg_scope(weight_decay=weight_decay)):
		net, end_points = resnet50(
			image_batch,
			num_classes=2,
			reuse=reuse,
			is_training=is_training
			)
	logits = tf.squeeze(net)
	# _activation_summaries(end_points)
	return logits, end_points

def loss(logits, label_batch, end_points):
	label_batch = tf.one_hot(label_batch,2)
	if 'AuxLogits' in end_points:
		tf.losses.softmax_cross_entropy(
			logits=end_points['AuxLogits'], onehot_labels=label_batch,
			label_smoothing=0.0, weights=0.4, scope='aux_loss')
	tf.losses.softmax_cross_entropy(
		logits=logits, onehot_labels=label_batch,
		label_smoothing=0.0, weights=1.0, scope='ini_loss')

